"""
Author: Jian SuiYue kaohx0214@yeah.net
Date: 2025-04-23 01:35:16
LastEditTime: 2025-05-15 17:12:45
FilePath: LibraryManagerStudio/database/db_manager.py
"""
import contextlib
import threading
from datetime import datetime
from pathlib import Path
from typing import Optional, Generator

from loguru import logger
from sqlalchemy import exc, inspect, text
from sqlmodel import create_engine, SQLModel, Session, select, and_

from .schema_manger import SchemaManger
from ..datamodel import Response


class DBManager:
    """
    DatabaseManager
    """
    __init_lock = threading.Lock()

    def __init__(self, engine_uri: str, base_dir: Optional[Path] = None):
        """
        Initialize the database manager with database connection settings.
        
        Args:
            engine_uri: Database connection URI (e.g. sqlite:///db.sqlite3) 
            base_dir: Base directory for migration files. If None, uses current directory
        """
        connection_args = {"check_same_thread": True} if "sqlite" in engine_uri else {}

        self.engine = create_engine(engine_uri, connect_args=connection_args)
        self.schema_manager = SchemaManger(engine=self.engine, base_dir=base_dir)
        self._session_factory = lambda: Session(self.engine)
        self._session = None

    def _should_auto_upgrade(self) -> bool:
        """
        Check if auto upgrade should run based on schema differences
        """
        needs_upgrade, _ = self.schema_manager.check_schema_status()
        return needs_upgrade

    def initialize_database(self, auto_upgrade: bool = False, force_init_alembic: bool = True) -> Response:
        """
        Initialize database and migration in the correct order.
        
        Args:
            auto_upgrade: If True, automatically generate and apply migrations for schema changes
            force_init_alembic: If True, reinitialize alembic configuration even if it exists
        """
        if not self.__init_lock.acquire(blocking=False):  # if not exist lock that is not a blocking lock 
            return Response(message="Database initialization already in progress", status=False)

        logger.info("Initializing database...")

        try:
            inspector = inspect(self.engine)
            tables_exist = inspector.get_table_names()  # pass inspector object to get all table name
            logger.info(f"Tables in the database: {tables_exist}")
            if not tables_exist:  # not exist table -> to create table
                logger.info("Creating database tables ...")
                SQLModel.metadata.create_all(self.engine)

                if self.schema_manager.initialize_migrations(force=force_init_alembic):
                    logger.info("Database initialized successfully")
                    return Response(message="Database initialized successfully", status=True)
                return Response(message="Failed to initialize migrations", status=False)

            # Handle existing database
            if auto_upgrade or self._should_auto_upgrade():
                logger.info("Checking database schema...")
                if self.schema_manager.ensure_schema_up_to_date():
                    logger.info("Database schema is up to date")
                    return Response(message="Database schema is up to date", status=True)
                return Response(message="Database upgrade failed", status=False)

            # logger.info("Database is ready")
            return Response(message="Database is ready", status=True)

        except Exception as e:
            error_msg = f'Database initialize failed: {str(e)}'
            logger.error(error_msg)
            return Response(message=error_msg, status=False)

        finally:
            self.__init_lock.release()

    async def verify_sql_include_table(self):
        """
        verify we create table success of not?
        """
        inspector = inspect(self.engine)
        tables_exist = inspector.get_table_names()
        if not tables_exist:
            error_msg = f'Verify database initialize failed: '
            logger.error(error_msg)
            return Response(message=error_msg, status=False)
        return Response(message="Verify database initialized successfully", status=True)

    def reset_db(self, recreate_table: bool = True):
        """
        Reset the database by dropping all tables and optionally creating them.
        """
        if not self.__init_lock.acquire(blocking=False):  # Can't get unblocking lock
            logger.warning("Database reset already in progress")
            return Response(message="Database reset already in progress", status=False, data=None)

        try:
            # Dispose existing connections
            self.engine.dispose()
            with Session(self.engine) as session:
                try:
                    # Disable foreign key checks for SQLite
                    if "sqlite" in str(self.engine.url):
                        session.exec(text("PRAGMA FOREIGN_KEYS=OFF"))  # sqlite3关闭键限制

                    # DROP all tables
                    SQLModel.metadata.drop_all(self.engine)
                    logger.info("ALl tables drop successfully")

                    if "sqlite" in str(self.engine.url):
                        session.exec(text("PRAGMA FOREIGN_KEYS=ON"))  # sqlite3开启键限制

                    session.commit()
                except Exception as e:
                    session.rollback()
                    raise e
                finally:
                    session.close()
                    self.__init_lock.release()

                if recreate_table:
                    logger.info("Recreating tables...")
                    self.initialize_database(auto_upgrade=True, force_init_alembic=True)

                return Response(
                    message="Database reset successfully" if recreate_table else "Database tables drop successfully",
                    status=True,
                    data=None
                )

        except Exception as e:
            error_msg = f"Error while resetting database: {str(e)}"
            logger.error(error_msg)
            return Response(message=error_msg, status=False, data=None)

        finally:
            if self.__init_lock.locked():
                self.__init_lock.release()
                logger.info("Database reset lock released")

    def upsert(self, model: SQLModel, return_json: bool = True):
        """Create or update an entity

        Args:
            model (SQLModel): The model instance to create or update
            return_json (bool, optional): If True, returns the model as a dictionary.
                If False, returns the SQLModel instance. Defaults to True.

        Returns:
            Response: Contains status, message and data (either dict or SQLModel based on return_json)
        """
        status = True
        model_class = type(model)
        existing_model = None

        with Session(self.engine) as session:
            try:
                existing_model = session.exec(select(model_class).where(model_class.id == model.id)).first()
                if existing_model:
                    if hasattr(model, "updated_at"):
                        model.updated_at = datetime.now()
                    for key, value in model.model_dump().items():
                        setattr(existing_model, key, value)
                    model = existing_model  # Use the updated existing model
                    session.add(model)
                else:
                    session.add(model)
                session.commit()
                session.refresh(model)
            except Exception as e:
                session.rollback()
                logger.error("Error while updating/creating " + str(model_class.__name__) + ": " + str(e))
                status = False

        return Response(
            message=(
                f"{model_class.__name__} Updated Successfully"
                if existing_model
                else f"{model_class.__name__} Created Successfully"
            ),
            status=status,
            data=model.model_dump() if return_json else model,
        )

    def _model_to_dict(self, model_obj):
        return {col.name: getattr(model_obj, col.name) for col in model_obj.__table__.columns}

    def get(
            self,
            model_class: SQLModel,
            filters: dict = None,
            page_info: dict = None,  # {pages: int, size: int}
            return_json: bool = False,
            order: str = "desc",
    ):
        """List entities"""
        with Session(self.engine) as session:
            result = []
            status = True
            status_message = ""
            pages = None

            try:
                statement = select(model_class)
                if filters:
                    conditions = [getattr(model_class, col) == value for col, value in filters.items()]
                    statement = statement.where(and_(*conditions))

                if page_info:
                    page, size = page_info['page'], page_info['size']
                    statement = statement.offset((page - 1) * size).limit(size)

                if hasattr(model_class, "created_at") and order:  # has attribute "created_at" -> need asc/desc by it
                    order_by_clause = getattr(model_class.created_at, order)()  # Dynamically apply asc/desc
                    statement = statement.order_by(order_by_clause)

                items = session.exec(statement)
                result = [self._model_to_dict(item) if return_json else item for item in items]
                status_message = f"{model_class.__name__} Retrieved Successfully"
            except Exception as e:
                session.rollback()
                status = False
                status_message = f"Error while fetching {model_class.__name__}"
                logger.error("Error while getting items: " + str(model_class.__name__) + " " + str(e))

            return Response(message=status_message, status=status, data=result)

    def delete(self, model_class: SQLModel, filters: dict = None):
        """Delete an entity"""
        status_message = ""
        status = True

        with Session(self.engine) as session:
            try:
                statement = select(model_class)
                if filters:
                    conditions = [getattr(model_class, col) == value for col, value in filters.items()]
                    statement = statement.where(and_(*conditions))

                rows = session.exec(statement).all()

                if rows:
                    for row in rows:
                        session.delete(row)
                    session.commit()
                    status_message = f"{model_class.__name__} Deleted Successfully"
                else:
                    status_message = "Row not found"
                    logger.info(f"Row with filters {filters} not found")

            except exc.IntegrityError as e:
                session.rollback()
                status = False
                status_message = (f"Integrity error: The {model_class.__name__} is linked to another entity and cannot "
                                  f"be deleted. {e}")
                # Log the specific integrity error
                logger.error(status_message)
            except Exception as e:
                session.rollback()
                status = False
                status_message = f"Error while deleting: {e}"
                logger.error(status_message)

        return Response(message=status_message, status=status, data=None)

    def get_session(self) -> Session:
        """
        get a session that use for contxet manager
        """
        return self._session_factory()

    @contextlib.contextmanager
    def session_scope(self) -> Generator[Session, None, None]:
        """
        context manager, to manager the session lifespan
        
        tutorial:
        with db_manager.session_scope() as session:
            # use session to do it
        """
        session = self._session_factory()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def close(self) -> None:
        """
        close the database connection
        """
        if self._session:
            self._session.close()
        self.engine.dispose()
        logger.info("Database connection closed")
